import requests
from bs4 import BeautifulSoup
import re
import json
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import os

def get_video_url(webpage_url):
    """
    从网页中提取视频源地址
    
    Args:
        webpage_url: 包含视频的网页地址
        
    Returns:
        视频源地址，如果未找到则返回None
    """
    # 设置Chrome选项
    chrome_options = Options()
    chrome_options.add_argument('--headless')  # 无头模式
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36')
    
    try:
        # 初始化WebDriver
        driver = webdriver.Chrome(options=chrome_options)
        print(f"正在访问网页: {webpage_url}")
        driver.get(webpage_url)
        
        # 等待页面加载完成
        time.sleep(3)
        
        # 等待视频元素出现
        try:
            WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.TAG_NAME, "video"))
            )
            print("找到视频元素")
        except Exception as e:
            print(f"等待视频元素超时: {e}")
        
        # 查找所有视频元素
        video_elements = driver.find_elements(By.TAG_NAME, "video")
        
        if not video_elements:
            print("未找到视频元素")
            driver.quit()
            return None
        
        # 提取视频源地址
        video_urls = []
        for video in video_elements:
            src = video.get_attribute('src')
            if src:
                video_urls.append(src)
            else:
                # 尝试从source标签获取
                source_elements = video.find_elements(By.TAG_NAME, "source")
                for source in source_elements:
                    src = source.get_attribute('src')
                    if src:
                        video_urls.append(src)
        
        # 关闭浏览器
        driver.quit()
        
        if video_urls:
            print(f"找到 {len(video_urls)} 个视频源地址")
            return video_urls
        else:
            print("未找到视频源地址")
            return None
            
    except Exception as e:
        print(f"获取视频源地址时出错: {e}")
        return None

def download_video(video_url, output_path):
    """
    下载视频到指定路径
    
    Args:
        video_url: 视频源地址
        output_path: 输出文件路径
    
    Returns:
        bool: 下载是否成功
    """
    try:
        print(f"正在下载视频: {video_url}")
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Referer': 'https://www.douyin.com/'
        }
        
        response = requests.get(video_url, headers=headers, stream=True)
        response.raise_for_status()
        
        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        
        print(f"视频已保存到: {output_path}")
        return True
    except Exception as e:
        print(f"下载视频时出错: {e}")
        return False

if __name__ == "__main__":
    # 获取用户输入的网页地址
    webpage_url = input("请输入包含视频的网页地址: ")
    
    # 获取视频源地址
    video_urls = get_video_url(webpage_url)
    print(f"视频地址: {video_urls[0]}")
    
    # if video_urls:
    #     # 创建输出目录
    #     output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "downloaded_videos")
    #     os.makedirs(output_dir, exist_ok=True)
        
    #     # 下载所有找到的视频
    #     for i, video_url in enumerate(video_urls):
    #         output_path = os.path.join(output_dir, f"video_{i+1}.mp4")
    #         download_video(video_url, output_path)
    # else:
    #     print("未能获取视频源地址，请检查网页地址是否正确")